import os
import numpy as np
import pandas as pd
from FSdata.FSdataset import FSdata, collate_fn, attr2catidx_map, idx2attr_map
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat
from FSdata.FSaug import  Resize, Compose,Normalize, UpperCrop
from utils.preprocessing import addFakeLabel
from utils.predicting import predict,gen_submission
import zipfile

class FSAugTest(object):
    def __init__(self):
        self.augment = Compose([
            Resize(size=(336,336),select=[0,1,2,5,6,7]),
            UpperCrop(size=(336,336), select=[3,4]),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, image, attr):
        return self.augment(image, attr)


os.environ["CUDA_VISIBLE_DEVICES"] = "2"

rawdata_root = '/media/gserver/data/FashionAI'

test_root = os.path.join(rawdata_root, 'rank')
test_pd = pd.read_csv(os.path.join(rawdata_root, 'rank/Tests/question.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])

test_pd = addFakeLabel(test_pd)

# select part
select_AttrIdx = [3,4]
select_AttrKey = [idx2attr_map[x] for x in select_AttrIdx]
test_pd = test_pd[test_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]


print test_pd.shape
print test_pd['AttrKey'].value_counts()

data_set = {}
data_set['test'] = FSdata(root_path=test_root,
                          anno_pd=test_pd,
                          transforms=FSAugTest(),
                          select=select_AttrIdx,
                           )

data_loader = {}
data_loader['test'] = torchdata.DataLoader(data_set['test'], batch_size=16, num_workers=4,
                                          shuffle=False, pin_memory=True,collate_fn=collate_fn)


# model prepare
model_name = 'resnet50_cat[3,4]_crop[3,4]-[0.9793]'
resume = '/media/gserver/models/FashionAI/resnet50_cat[3,4]_crop[3,4]/bestweights-[0.8810]-[0.9793].pth'
model = resnet50_cat(pretrained=True, num_classes=data_set['test'].num_classes)

model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
model = model.cuda()

# predict
if not os.path.exists('./online_pred/csv'):
    os.makedirs('./online_pred/csv')
if not os.path.exists('./online_pred/subs_zip'):
    os.makedirs('./online_pred/subs_zip')
if not os.path.exists('./online_pred/part_csv'):
    os.makedirs('./online_pred/part_csv')

# save csv and zip
pred_scores, _, _, _ = predict(model, data_set['test'], data_loader['test'])
test_pred = gen_submission(test_pd[['ImageName', 'AttrKey']], pred_scores, catidx_map=data_set['test'].catidx_map)

if len(select_AttrIdx) < 8:
    test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred/part_csv/%s.csv' % model_name,
                                                                 header=None, index=False)
else:
    test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred/csv/%s.csv' % model_name,
                                                                 header=None, index=False)
    # make zip file
    z = zipfile.ZipFile('./online_pred/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
    z.write('./online_pred/csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
    z.close()

print test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].info()
